﻿/*
// Exepack.NET 
// http://www.codeplex.com/exepack
//
// Executable compressor
// Written by Y [07-03-09]
// Copyright (c) 2008-2010 Alexey Yakovlev
*/

using System;
using System.IO;
using System.Text;
using System.Linq;
using System.IO.Compression;
using System.Collections.Generic;
using System.Runtime.Serialization;
using System.CodeDom.Compiler;
using Microsoft.CSharp;
using System.ComponentModel;
using System.Reflection.Emit;
using System.Threading;
using System.Reflection;
using BLToolkit;
using System.Security.Policy;
using Vestris.ResourceLib;

namespace Exepack
{
	using AssemblyFactory = Mono.Cecil.AssemblyFactory;

	/// <summary>
	/// Compresses .NET executables
	/// </summary>
	public class ExeCompressor : IDisposable
	{
		TempFileManager temp = new TempFileManager();

		public ExeCompressor()
		{ 
		}

		~ExeCompressor()
		{
			Dispose();
		}

		public void Dispose()
		{
			if (temp != null)
			{
				temp.Dispose();
			}

			temp = null;
		}

		/// <summary>
		/// Reports compression progress
		/// </summary>
		public event EventHandler<CompressionProgressEventArgs> CompressionProgress;

		/// <summary>
		/// Generating compressed executable progress
		/// </summary>
		public event EventHandler<CompressionProgressEventArgs> SavingExecutableProgress;

		/// <summary>
		/// Compress the specified file using default compression options
		/// </summary>
		public void CompressExecutable(string fileName)
		{
			CompressExecutable(CompressionOptions.Create(fileName));
		}

		/// <summary>
		/// Compress the specified file using provided compression options
		/// </summary>
		public void CompressExecutable(CompressionOptions options)
		{
			// list of resources
			var resList = new List<string>();

			// compress main assembly
			var asm = options.MainAssembly;
			resList.Add(asm.ResourceName);
			CompressFile(asm.FullPath, temp[asm.ResourceName]);

			// compress referenced assemblies
			foreach (var addAsm in options.AdditionalAssemblies.Where(
				v => v.IsILOnly && !v.IsSystemAssembly))
			{
				resList.Add(addAsm.ResourceName);
				CompressFile(addAsm.FullPath, temp[addAsm.ResourceName]);
			}

			var inFileSize = options.UncompressedSize;
			SavingExecutableProgress.Trigger(this, new CompressionProgressEventArgs(options.OutputFileName, inFileSize, 90));

			// save compressed executable
			GenerateLoader(options, resList.ToArray());
			OptimizeAssembly(options.OutputFileName);
			CopyResources(options);

			// display statistics
			var outFileSize = new FileInfo(options.OutputFileName).Length;
			SavingExecutableProgress.Trigger(this, new CompressionProgressEventArgs(options.OutputFileName, inFileSize, (int)outFileSize, 100));
		}

		/// <summary>
		/// Deflate single file, report progress
		/// </summary>
		void CompressFile(string inputFile, string outputFile)
		{
			var fileName = Path.GetFileName(inputFile);
			var length = 0;

			using (var inFile = File.OpenRead(inputFile))
			using (var outFile = File.Create(outputFile))
			using (var headerWriter = new BinaryWriter(outFile))
			{
				length = (int)inFile.Length;
				CompressionProgress.Trigger(this, new CompressionProgressEventArgs(fileName, length, 0));
				headerWriter.Write(length);

				using (var ds = new DeflateStream(outFile, CompressionMode.Compress))
				using (var writer = new BinaryWriter(ds))
				using (var reader = new BinaryReader(inFile))
				{
					writer.Write(reader.ReadBytes(length));
				}
			}

			var packed = new FileInfo(outputFile).Length;
			CompressionProgress.Trigger(this, new CompressionProgressEventArgs(fileName, length, (int)packed, 50));
		}

		/// <summary>
		/// Copy win32 resources (only icons and version info)
		/// </summary>
		private void CopyResources(CompressionOptions options)
		{
			if (!options.CopyIcons && !options.CopyVersionInfo)
				return;

			// copy icon and version resources from one PE to another as needed
			using (var vi = new ResourceInfo())
			{
				vi.Load(options.MainAssembly.FullPath);

				foreach (var type in vi.ResourceTypes.Where(t =>
					(options.CopyIcons && t.ResourceType == Kernel32.ResourceTypes.RT_GROUP_ICON) ||
					(options.CopyVersionInfo && t.ResourceType == Kernel32.ResourceTypes.RT_VERSION)))
				{
					foreach (var res in vi.Resources[type])
					{
						res.SaveTo(options.OutputFileName);
					}
				}
			}
		}

		/// <summary>
		/// Optimize generated assembly for free :)
		/// Just load and re-save it using Mono.Cecil library
		/// </summary>
		void OptimizeAssembly(string fileName)
		{
			var myLibrary = AssemblyFactory.GetAssembly(fileName);
			AssemblyFactory.SaveAssembly(myLibrary, fileName);
		}

		/// <summary>
		/// Creates random name for compressed file loader
		/// </summary>
		internal string GenerateRandomName()
		{
			return Convert.ToBase64String(Guid.NewGuid().ToByteArray());
		}

		/// <summary>
		/// Generate loader for compressed executable.
		/// The code is the same as in Loader class.
		/// </summary>
		internal void GenerateLoader(CompressionOptions options, string[] resources)
		{
			// create a dynamic assembly and module 
			var assemblyName = new AssemblyName(options.OutputAssemblyName);
			var assemblyBuilder = Thread.GetDomain().DefineDynamicAssembly(assemblyName, AssemblyBuilderAccess.Save);
			var module = assemblyBuilder.DefineDynamicModule(assemblyName.Name, options.OutputFileName);

			// add compressed assemblies as private resources
			foreach (var resource in resources)
			{
				module.DefineManifestResource(resource, File.OpenRead(temp[resource]), ResourceAttributes.Private);
			}

			// int Main(string[] args)
			var mainBuilder = module.DefineGlobalMethod(GenerateRandomName(), // "Main"
				MethodAttributes.HideBySig | MethodAttributes.Static | MethodAttributes.Public,
				typeof(int), new Type[] { typeof(string[]) });

			// [STAThread] attribute
			mainBuilder.SetCustomAttribute(new CustomAttributeBuilder(
				typeof(STAThreadAttribute).GetConstructor(new Type[0]), new object[0]));

			// Assembly ExtractAssembly(object sender, ResolveEventArgs args)
			var extractAssemblyBuilder = module.DefineGlobalMethod(GenerateRandomName(), // "ExtractAssembly"
				MethodAttributes.HideBySig | MethodAttributes.Static | MethodAttributes.Public,
				typeof(Assembly), new Type[] { typeof(object), typeof(ResolveEventArgs) });

			// Assembly GetResourceAssembly(string name)
			var getResourceAssemblyBuilder = module.DefineGlobalMethod(GenerateRandomName(), // "GetResourceAssembly"
				MethodAttributes.HideBySig | MethodAttributes.Static | MethodAttributes.Public,
				typeof(Assembly), new Type[] { typeof(string) });

			// int Main(string[] args)
			var asm = new EmitHelper(mainBuilder.GetILGenerator())

				// AppDomain.CurrentDomain.AssemblyResolve += ExtractAssembly;
				.call(typeof(AppDomain).GetMethod("get_CurrentDomain"))
				.dup
				.ldnull
				.ldftn(extractAssemblyBuilder)
				.newobj(typeof(ResolveEventHandler).GetConstructor(new Type[] { typeof(object), typeof(IntPtr) }))
				.callvirt(typeof(AppDomain).GetMethod("add_AssemblyResolve", new Type[] { typeof(ResolveEventHandler) }))

				// return AppDomain.CurrentDomain.ExecuteAssemblyByName(new AssemblyName("Assembly, Version=0.0.0.0, ..."), null, args);
				//.call(typeof(AppDomain).GetMethod("get_CurrentDomain")) // dup already did it
				.ldstr(options.MainAssembly.FullName.ToLowerInvariant())
				.newobj(typeof(AssemblyName).GetConstructor(new Type[] { typeof(string) }))
				.ldnull
				.ldarg(0)
				.callvirt(typeof(AppDomain).GetMethod("ExecuteAssemblyByName", new Type[] { typeof(AssemblyName), typeof(Evidence), typeof(string[]) }))
				.ret();

			// Assembly ExtractAssembly(object sender, ResolveEventArgs args)
			asm = new EmitHelper(extractAssemblyBuilder.GetILGenerator())

				// return GetResourceAssembly(Convert.ToBase64String(Encoding.Default.GetBytes(args.Name.ToLowerInvariant())));
				.call(typeof(Encoding).GetMethod("get_Default"))
				.ldarg(1)
				.callvirt(typeof(ResolveEventArgs).GetMethod("get_Name"))
				.callvirt(typeof(String).GetMethod("ToLowerInvariant"))
				.callvirt(typeof(Encoding).GetMethod("GetBytes", new Type[] { typeof(string) }))
				.call(typeof(Convert).GetMethod("ToBase64String", new Type[] { typeof(byte[]) }))
				.call(getResourceAssemblyBuilder)
				.ret();

			// Assembly GetResourceAssembly(string name)
			asm = new EmitHelper(getResourceAssemblyBuilder.GetILGenerator());
			var streamVar = asm.DeclareLocal(typeof(Stream));
			var lengthVar = asm.DeclareLocal(typeof(int));
			var defStreamVar = asm.DeclareLocal(typeof(DeflateStream));
			var assemblyVar = asm.DeclareLocal(typeof(Assembly));
			var streamNotNullLabel = asm.DefineLabel();
			var streamNullLabel = asm.DefineLabel();
			var defStreamNullLabel = asm.DefineLabel();
			var endTryBlockLabel = asm.DefineLabel();

			asm
				// var stream = Assembly.GetExecutingAssembly().GetManifestResourceStream(name)
				.call(typeof(Assembly).GetMethod("GetExecutingAssembly"))
				.ldarg(0)
				.callvirt(typeof(Assembly).GetMethod("GetManifestResourceStream", new Type[] { typeof(string) }))
				.stloc(streamVar)

				// using (stream)
				.BeginExceptionBlock();
			asm
				// if (stream == null) return null;
				.ldloc(streamVar)
				.dup
				.brtrue(streamNotNullLabel)
				.pop
				.ldnull
				.stloc(assemblyVar)
				.leave(endTryBlockLabel)
				.MarkLabel(streamNotNullLabel)
				
				// int length = new BinaryReader(stream).ReadInt32();
				.newobj(typeof(BinaryReader).GetConstructor(new Type[] { typeof(Stream) })) // trick Reflector: remember .dup?
				.callvirt(typeof(BinaryReader).GetMethod("ReadInt32"))
				.stloc(lengthVar)

				// var defStream = new DeflateStream(stream, CompressionMode.Decompress)
				.ldloc(streamVar)
				.ldc_i4(0)
				.newobj(typeof(DeflateStream).GetConstructor(new Type[] { typeof(Stream), typeof(CompressionMode) }))
				.stloc(defStreamVar)

				// using (defStream)
				.BeginExceptionBlock();
			asm
				// var assembly = Assembly.Load(new BinaryReader(ds).ReadBytes(length));
				.ldloc(defStreamVar)
				.newobj(typeof(BinaryReader).GetConstructor(new Type[] { typeof(Stream) }))
				.ldloc(lengthVar)
				.callvirt(typeof(BinaryReader).GetMethod("ReadBytes", new Type[] { typeof(int) }))
				.call(typeof(Assembly).GetMethod("Load", new Type[] { typeof(byte[]) }))
				.stloc(assemblyVar)
				.leave_s(endTryBlockLabel)

			.BeginFinallyBlock()

				// if (defStream != null)
				.ldloc(defStreamVar)
				.ldnull
				.ceq
				.brtrue_s(defStreamNullLabel)

				// defStream.Dispose()
				.ldloc(defStreamVar)
				.callvirt(typeof(IDisposable).GetMethod("Dispose"))
				.MarkLabel(defStreamNullLabel)
				.EndExceptionBlock()

			.BeginFinallyBlock()

				// if (stream != null)
				.ldloc(streamVar)
				.ldnull
				.ceq
				.brtrue_s(streamNullLabel)

				// stream.Dispose()
				.ldloc(streamVar)
				.callvirt(typeof(IDisposable).GetMethod("Dispose"))
				.MarkLabel(streamNullLabel)
				.EndExceptionBlock()

			// return assembly
			.MarkLabel(endTryBlockLabel)
			.ldloc(assemblyVar)
			.ret();

			// bake <Module> type
			module.CreateGlobalFunctions();

			// set the entry point for the application and save it
			assemblyBuilder.SetEntryPoint(mainBuilder, 
				options.MainAssembly.IsConsoleApplication ? PEFileKinds.ConsoleApplication : PEFileKinds.WindowApplication);
			assemblyBuilder.Save(options.OutputFileName);
		}
	}
}
